#include "VulkanApplication.h"
#include "BufferManager.h"
#include "ufbx.h"
#include "../glm/glm.hpp"
#include "../glm/gtc/matrix_transform.hpp"
#include "../glm/gtc/quaternion.hpp"
#include <vector>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

struct DrawObject
{
    std::string descriptorName;
    std::string textureName;
    std::string vertexBufferName;
    std::string indexBufferName;
    unsigned int indexCount;
};

glm::vec3 cam_pos(2.0f, 0.0f, 0.0f), cam_target(0.0f, 0.0f, 0.0f), cam_up(0.0f, 1.0f, 0.0f);

void VulkanApplication::keyCallback(GLFWwindow* window, int key, int scancode, int action, int mods)
{
    if (action > 0)
    {
        glm::vec3 forward = cam_target - cam_pos;
        float distance = glm::length(forward);
        forward = glm::normalize(forward);
        
        glm::vec3 side = glm::cross(cam_up, forward);
        if (key == GLFW_KEY_A)
        {
            forward = glm::angleAxis(-0.1f, cam_up) * forward;
            cam_pos = cam_target - forward * distance;
        }
        else if (key == GLFW_KEY_D)
        {
            forward = glm::angleAxis(0.1f, cam_up) * forward;
            cam_pos = cam_target - forward * distance;
        }
        else if (key == GLFW_KEY_S)
        {
            forward = glm::angleAxis(-0.1f, side) * forward;
            cam_pos = cam_target - forward * distance;
        }
        else if (key == GLFW_KEY_W)
        {
            forward = glm::angleAxis(0.1f, side) * forward;
            cam_pos = cam_target - forward * distance;
        }
        else if (key == GLFW_KEY_Z)
        {
            distance *= 0.95f;
            cam_pos = cam_target - forward * distance;
        }
        else if (key == GLFW_KEY_Q)
        {
            distance *= 1.05f;
            cam_pos = cam_target - forward * distance;
        }
        else if (key == GLFW_KEY_C)
        {
            cam_pos.y -= 0.1f; cam_target.y -= 0.1f;
        }
        else if (key == GLFW_KEY_E)
        {
            cam_pos.y += 0.1f; cam_target.y += 0.1f;
        }
        cam_up = glm::cross(forward, side);
    }
    if (key == GLFW_KEY_ESCAPE) glfwSetWindowShouldClose(window, 1);
}

std::vector<DrawObject> loadModel(VulkanApplication& app, BufferManager& manager, const std::string& file,
                                  glm::vec3& minPoint, glm::vec3& maxPoint)
{
    ufbx_load_opts opts = { 0 }; ufbx_error error;
    ufbx_scene* scene = ufbx_load_file(file.c_str(), &opts, &error);
    if (!scene)
    {
        printf("Failed to load: %s\n", error.description.data);
        return std::vector<DrawObject>();
    }

    minPoint = glm::vec3(9999.0f, 9999.0f, 9999.0f);
    maxPoint = glm::vec3(-9999.0f, -9999.0f, -9999.0f);

    std::vector<DrawObject> drawObjects;
    for (size_t i = 0; i < scene->nodes.count; ++i)
    {
        ufbx_node* node = scene->nodes.data[i];
        if (!node->mesh) continue;

        ufbx_mesh* mesh = node->mesh;
        size_t maxFaceIndices = mesh->max_face_triangles * 3;
        for (size_t n = 0; n < mesh->material_parts.count; ++n)
        {
            std::vector<Vertex> vertices;
            ufbx_vec4 defColor; defColor.x = 1.0f; defColor.y = 1.0f; defColor.z = 1.0f;
            ufbx_vec2 defUV; defUV.x = 0.0f, defUV.y = 0.0f;
            size_t idCount = 0;

            ufbx_mesh_part meshPart = mesh->material_parts[n];
            ufbx_material* material = mesh->materials[meshPart.index];
            unsigned int* indices = (unsigned int*)calloc(maxFaceIndices, sizeof(unsigned int));
            for (size_t f = 0; f < meshPart.num_faces; ++f)
            {
                ufbx_face face = mesh->faces.data[meshPart.face_indices.data[f]];
                unsigned int numTriangles = ufbx_triangulate_face(indices, maxFaceIndices, mesh, face);
                for (size_t k = 0; k < numTriangles * 3; ++k)
                {
                    unsigned int idx = indices[k];
                    ufbx_vec3 pos = ufbx_get_vertex_vec3(&mesh->vertex_position, idx);
                    ufbx_vec4 color = mesh->vertex_color.exists
                                    ? ufbx_get_vertex_vec4(&mesh->vertex_color, idx) : defColor;
                    ufbx_vec2 uv = mesh->vertex_uv.exists
                                 ? ufbx_get_vertex_vec2(&mesh->vertex_uv, idx) : defUV;
                    idCount++;

                    Vertex vec;
                    vec.x = (float)pos.x; vec.y = (float)pos.y; vec.z = (float)pos.z;
                    vec.r = (float)color.x; vec.g = (float)color.y; vec.b = (float)color.z;
                    vec.u = (float)uv.x; vec.v = 1.0f - (float)uv.y;

                    if (vec.x < minPoint.x) minPoint.x = vec.x; if (vec.x > maxPoint.x) maxPoint.x = vec.x;
                    if (vec.y < minPoint.y) minPoint.y = vec.y; if (vec.y > maxPoint.y) maxPoint.y = vec.y;
                    if (vec.z < minPoint.z) minPoint.z = vec.z; if (vec.z > maxPoint.z) maxPoint.z = vec.z;
                    vertices.push_back(vec);
                }
            }
            free(indices);

            ufbx_vertex_stream dataStream[1];
            dataStream[0].data = &vertices[0];
            dataStream[0].vertex_size = sizeof(Vertex);
            dataStream[0].vertex_count = vertices.size();

            ufbx_error error;
            std::vector<unsigned int> triangles(meshPart.num_triangles * 3);
            size_t vCount = ufbx_generate_indices(dataStream, 1, &triangles[0], idCount, NULL, &error);
            if (error.type != UFBX_ERROR_NONE)
            {
                printf("Failed to generate indices: %s\n", error.description.data);
                continue;
            }
            vertices.resize(vCount);

            std::string name = mesh->name.data ? std::string(mesh->name.data) + std::to_string(n)
                             : "mesh" + std::to_string(i) + "_" + std::to_string(n);
            manager.createVertexBuffer(name + "_v", vertices);
            manager.createIndexBuffer(name + "_i", triangles);
            if (material && material->textures.data)
            {
                ufbx_texture* tex = material->textures[0].texture;
                VkImage image = nullptr;
                if (tex && tex->content.data)
                {
                    image = manager.createTextureImage(
                        name + "_t", (unsigned char*)tex->content.data, (int)tex->content.size);
                }
                else if (tex && tex->filename.data)
                {
                    image = manager.createTextureImage(
                        name + "_t", std::string(tex->filename.data));
                }

                if (!image)
                {
                    static unsigned char defaultBMP[] = {
                        0x42, 0x4d, 0x3a, 0, 0, 0, 0, 0, 0, 0, 0x36, 0, 0, 0, 0x28,
                        0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0x18, 0, 0, 0, 0, 0,
                        4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                        0xff, 0xff, 0xff, 0 };
                    image = manager.createTextureImage(name + "_t", defaultBMP, 58);
                }

                manager.createTextureSampler(name + "_t", image, VK_FORMAT_R8G8B8A8_SRGB);
                manager.allocateNewDescriptor(name + "_d");
                manager.updateTextureDescriptor(name + "_d", name + "_t",
                                                VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
            }

            DrawObject drawObject
            {
                name + "_d", name + "_t", name + "_v",
                name + "_i", (unsigned int)triangles.size()
            };
            drawObjects.push_back(drawObject);
        }
    }

    ufbx_free_scene(scene);
    return drawObjects;
}

int main(int argc, char** argv)
{
    VulkanApplication app("c16_model", 1024, 768);
    if (!app.initialized())
    {
        printf("Failed to initialize VulkanApplication.\n");
        return 1;
    }

    // Create pipeline
    BufferManager manager(&app);
    if (!manager.createDescriptorSet() ||
        !manager.prepare("draw_texture.vert", "draw_texture.frag"))
    {
        printf("Failed to create pipeline.\n");
        return 1;
    }

    // Load model data and convert them to "draw objects"
    glm::vec3 minPoint, maxPoint;
    std::string model((argc > 1) ? argv[1] : "CesiumMan.fbx");

    std::vector<DrawObject> drawObjects = loadModel(app, manager, model, minPoint, maxPoint);
    cam_target = (minPoint + maxPoint) * 0.5f;
    cam_pos = cam_target + glm::vec3((maxPoint.x - minPoint.x) * 5.0f, 0.0f, 0.0f);

    // Define UBO structure and apply to descriptors
    struct UniformBufferObject
    {
        glm::mat4x4 modelview;
        glm::mat4x4 projection;
    };
    size_t uboSize = sizeof(UniformBufferObject);
    UniformBufferObject ubo;

    VkBuffer uboBuffer = manager.createUniformBuffer("ubo", uboSize);
    void* uboData = manager.mapBufferMemory("ubo", uboSize);
    for (size_t i = 0; i < drawObjects.size(); ++i)
        manager.updateUniformDescriptor(drawObjects[i].descriptorName, "ubo", uboSize);
    
    // Main loop
    int width = 1024, height = 768;
    VkFence waitFence = app.getOrCreateFence("CommandFence");
    VkSemaphore waitNextImage = app.getOrCreateSemaphore("WaitNextImage");
    VkSemaphore waitSubmission = app.getOrCreateSemaphore("WaitSubmission");
    VkCommandBuffer commandBuffer = app.getCommandBuffer();

    float angleZ = 0.0f, aspectRatio = 0.0f;
    while (!app.shouldClose())
    {
        vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
        vkResetFences(app.getDevice(), 1, &waitFence);

        uint32_t imageIndex = 0;
        VkResult result = vkAcquireNextImageKHR(app.getDevice(), app.getSwapchain(), UINT64_MAX,
            waitNextImage, VK_NULL_HANDLE, &imageIndex);
        switch (result)
        {
        case VK_SUBOPTIMAL_KHR:
        case VK_ERROR_OUT_OF_DATE_KHR:
            imageIndex = app.recreateSwapchain(width, height, waitNextImage); break;
        case VK_SUCCESS:
            break;
        default:
            printf("Next image: %d (RESULT = %d)\n", imageIndex, result);
            app.setClosed(); break;
        }

        aspectRatio = (float)width / (float)height;
        ubo.modelview = glm::lookAt(cam_pos, cam_target, cam_up);
        ubo.projection = glm::perspective(45.0f, aspectRatio, 0.1f, 1000.0f);
        memcpy(uboData, &ubo, sizeof(UniformBufferObject));

        app.beginRendering(imageIndex, width, height);
        for (size_t i = 0; i < drawObjects.size(); ++i)
        {
            DrawObject& drawObject = drawObjects[i];
            VkBuffer vertexBuffer = manager.getBuffer(drawObject.vertexBufferName);
            VkBuffer indexBuffer = manager.getBuffer(drawObject.indexBufferName);

            VkDescriptorSet descriptorSet = manager.getDescriptorSet(drawObject.descriptorName);
            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                app.getPipelineLayout(), 0, 1, &descriptorSet, 0, nullptr);

            VkBuffer vertexBuffers[] = { vertexBuffer };
            VkDeviceSize offsets[] = { 0 };
            vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
            vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
            vkCmdDrawIndexed(commandBuffer, drawObject.indexCount, 1, 0, 0, 0);
        }
        app.endRendering();
        app.submitAndPresent(waitNextImage, waitSubmission, waitFence, imageIndex);
        glfwPollEvents();
    }

    vkDeviceWaitIdle(app.getDevice());
    manager.unmapBufferMemory("ubo");
    manager.cleanup();
    return 0;
}